from Function import *
from SparseMatrix import *

m = 512
n = 1024

def mu_lasso(vec,delta,mu):
    ans = 0
    for i in range(vec.shape[1]):
        x_i = vec[i,0]
        if abs(x_i) < delta:
            ans += x_i*2/(2*delta)
        else:
            ans += abs(x_i)-delta/2
    return mu*ans
        
A = np.random.normal(size = (m,n))
b = np.random.normal(size = (m,1))

class LASSO1(Function):
    def f(self,vec):
        mu = 1e-2
        delta = 1e-3*mu
        return norm(np.matmul(A,vec)-b)**2/2 + mu_lasso(vec,delta,mu)
    
class LASSO2(Function):
    def f(self,vec):
        mu = 1e-3
        delta = 1e-3*mu
        return norm(np.matmul(A,vec)-b)**2/2 + mu_lasso(vec,delta,mu)

F1 = LASSO1()
F2 = LASSO2()
x0 = sprase_rand(n,1,0.1)
print("For mu = 1e-2:")
ans = Steepest_Descent(F1,x0,epsilon = 1e-2)
print("Ans:")
print(ans)
print("with error:",norm(grad(F1,ans)))
print("For mu = 1e-3:")
ans = Steepest_Descent(F2,x0,epsilon = 1e-2)
print("Ans:")
print(ans)
print("with error:",norm(grad(F2,ans)))
print("For mu = 1e-2:")
ans = Barzilai_Borwein(F1,x0,epsilon = 1e-2)
print("Ans:")
print(ans)
print("with error:",norm(grad(F1,ans)))
print("For mu = 1e-3:")
ans = Barzilai_Borwein(F2,x0,epsilon = 1e-2)
print("Ans:")
print(ans)
print("with error:",norm(grad(F2,ans)))
